# -*- coding: utf-8 -*-
"""
Created on 04/05/2023
@author: 曾导SJTU
"""

# 迭代求解

import numpy as np
from precision_data import fp

def eqn_scalar_norm2(case, dim, it_nl, ncx, ncy, ncz, u0, u, var):

    # 定义与case对象相关的局部变量
    if var == "temp":
        l2_u = case.l2_T
        l2_max_u = case.l2_max_T

    # 计算L2范数
    
    # 等价形式
    # for k in range(ncz):
    #     for j in range(ncy):
    #         for i in range(ncx):
    #             res_u = u[i, j, k] - u0[i, j, k]
    #             l2_u += res_u**2
    # ncell = ncx*ncy*ncz
    # l2_u = math.sqrt(l2_u/fp(ncell))
    
    l2_u = np.sqrt(np.mean((u - u0)**2))

    # 计算到目前为止观察到的最大L2范数
    l2_max_u = max(l2_u, l2_max_u)

    return (l2_u, l2_max_u)

# 雅克比迭代
def scalar_pj(case, post, dim, it_nl, niter, relax, ncx, ncy, ncz, ncoef, cT, T, initzero, res):
    
    # 定义与case对象相关的局部变量
    id_aP = case.id_aP
    id_aE = case.id_aE
    id_aW = case.id_aW
    id_aN = case.id_aN
    id_aS = case.id_aS
    if dim == 3:
        id_aT = case.id_aT
        id_aB = case.id_aB
    id_bsrc = case.id_bsrc
    
    nsteps = case.nsteps
    
    # 定义与post对象相关的局部变量
    res_freq = post.res_freq
    linsol_fname = post.linsol_fname
    
    if initzero:
        T.fill(fp(0.0))

    # 分配一个数组并初始化它
    T0 = np.zeros((ncx, ncy, ncz), dtype=fp)

    max_norm = -1.e20

    # 开始线性求解器迭代    niter温度线性求解器迭代次数
    for it in range(1, niter+1):
        
        # 将当前版本复制到旧版本
        T0 = T.copy()

        # 初始化标准
        norm2 = 0.0

        # 对每个网格都循环  在每个网格迭代的过程中都是与周围网格的数值有关系，下面这个就是在找到每个网格周围的网格上的数据
        for k in range(ncz):
            for j in range(ncy):
                for i in range(ncx):
                    
                    if i == 0:
                        Tw = 0.0
                    else:
                        Tw = T0[i-1,j,k]
                    if i == ncx-1:
                        Te = 0.0
                    else:
                        Te = T0[i+1,j,k]
                    
                    if j == 0:
                        Ts = 0.0
                    else:
                        Ts = T0[i,j-1,k]
                    if j == ncy-1:
                        Tn = 0.0
                    else:
                        Tn = T0[i,j+1,k]
                    
                    if dim == 3:
                        if k == 0:
                            Tb = 0.0
                        else:
                            Tb = T0[i,j,k-1]
                        if k == ncz-1:
                            Tt = 0.0
                        else:
                            Tt = T0[i,j,k+1]

                    # \这个代表连接符号，以下代码是一行的意思，这个就是迭代公式，出现在Lecture 1.35，注意，T_new = T_new / cT[i,j,k,id_aP]这样才是完整的迭代.
                    T_new = (                 \
                        - cT[i,j,k,id_aE]*Te  \
                        - cT[i,j,k,id_aW]*Tw  \
                        - cT[i,j,k,id_aN]*Tn  \
                        - cT[i,j,k,id_aS]*Ts  \
                        + cT[i,j,k,id_bsrc]   \
                            )

                    if dim==3:
                        T_new = T_new - cT[i,j,k,id_aT]*Tt - cT[i,j,k,id_aB]*Tb

                    T_new = T_new / cT[i,j,k,id_aP]

                    du = relax * (T_new - T0[i,j,k]) # relax
                    # 更新t温度场
                    T[i,j,k] = T0[i,j,k] + du
                    norm2 = norm2 + du*du

        # 网格数量
        ncell = ncx*ncy*ncz
        # 计算L2范数
        norm2 = np.sqrt(norm2 / fp(ncell))

        if it == 1:     # it是迭代次数 
            initial_norm = norm2+1.e-20

        max_norm = max(norm2, max_norm) + 1.e-20

        rel_norm = norm2 / max_norm

        if rel_norm < res or it == niter:   # res 线性求解器迭代的收敛误差判定值    niter 线性求解器迭代的最大次数
            case.total_linsol_iters = case.total_linsol_iters + it
            if it_nl % res_freq == 0 or it_nl == 1 or it_nl == nsteps:
                print('it_nl, it, tot_it, norm2, init, max, rel ', it_nl, it, case.total_linsol_iters, norm2, initial_norm, max_norm, rel_norm)
                with open(linsol_fname, 'a') as linsol_fid:
                    print(it_nl, it, case.total_linsol_iters, norm2, initial_norm, max_norm, rel_norm, file=linsol_fid)
            break

    return

# 高斯赛戴尔迭代
def scalar_gs(case, post, dim, it_nl, niter, relax, ncx, ncy, ncz, ncoef, cT, T, initzero, res):
    
    # define local variables related to the case object
    id_aP = case.id_aP
    id_aE = case.id_aE
    id_aW = case.id_aW
    id_aN = case.id_aN
    id_aS = case.id_aS
    if dim == 3:
        id_aT = case.id_aT
        id_aB = case.id_aB
    id_bsrc = case.id_bsrc
    
    nsteps = case.nsteps
    
    # define local variables related to the post object
    res_freq = post.res_freq
    linsol_fname = post.linsol_fname
    
    if initzero:
        T.fill(fp(0.0))

    max_norm = -1.e20

    # Begin linear solver iterations
    for it in range(1, niter+1):
        
        # Initialize norm
        norm2 = 0.0

        for k in range(ncz):
            for j in range(ncy):
                for i in range(ncx):
                    
                    if i == 0:
                        Tw = 0.0
                    else:
                        Tw = T[i-1,j,k]
                    if i == ncx-1:
                        Te = 0.0
                    else:
                        Te = T[i+1,j,k]
                    
                    if j == 0:
                        Ts = 0.0
                    else:
                        Ts = T[i,j-1,k]
                    if j == ncy-1:
                        Tn = 0.0
                    else:
                        Tn = T[i,j+1,k]
                    
                    if dim == 3:
                        if k == 0:
                            Tb = 0.0
                        else:
                            Tb = T[i,j,k-1]
                        if k == ncz-1:
                            Tt = 0.0
                        else:
                            Tt = T[i,j,k+1]

                    T_new = (                 \
                        - cT[i,j,k,id_aE]*Te  \
                        - cT[i,j,k,id_aW]*Tw  \
                        - cT[i,j,k,id_aN]*Tn  \
                        - cT[i,j,k,id_aS]*Ts  \
                        + cT[i,j,k,id_bsrc]   \
                            )

                    if dim==3:
                        T_new = T_new - cT[i,j,k,id_aT]*Tt - cT[i,j,k,id_aB]*Tb

                    T_new = T_new / cT[i,j,k,id_aP]

                    du = relax * (T_new - T[i,j,k])
                    T[i,j,k] = T[i,j,k] + du
                    norm2 = norm2 + du*du

        ncell = ncx*ncy*ncz
        norm2 = np.sqrt(norm2 / fp(ncell))

        if it == 1:
            initial_norm = norm2+1.e-20

        max_norm = max(norm2, max_norm) + 1.e-20

        rel_norm = norm2 / max_norm

        if rel_norm < res or it == niter:
            case.total_linsol_iters = case.total_linsol_iters + it
            if it_nl % res_freq == 0 or it_nl == 1 or it_nl == nsteps:
                print('it_nl, it, tot_it, norm2, init, max, rel ', it_nl, it, case.total_linsol_iters, norm2, initial_norm, max_norm, rel_norm)
                with open(linsol_fname, 'a') as linsol_fid:
                    print(it_nl, it, case.total_linsol_iters, norm2, initial_norm, max_norm, rel_norm, file=linsol_fid)
            break

    return

#    def scalar_adi ...
